#!/bin/bash
# Copyright (C) 2024-2025 KoudaiPC Computer Technologies Co., Ltd.
# script name: yandi-host-build-nocloud-arm64-image.sh
# author: david921518@qq.com

################
# global vars
GLB_WORK_DIR=
GLB_DL_IMG_FILE=debian-12-nocloud-arm64-20251112-2294.qcow2
GLB_DL_IMG_MD5="e5b4bb29cd514f4b5dceecfc5531739f"
GLB_DL_URL=https://cloud.debian.org/images/cloud/bookworm/20251112-2294/$GLB_DL_IMG_FILE
################

user=$(whoami)

# 定义日志格式-前缀
log_prefix="[yandi-build]"
# 定义日志目录和文件名
log_file_dir="/var/log/kpc/yandi"
log_file_path="$log_file_dir/host-build.log"

log_kmsg_path="/dev/kmsg"
log_con_path="/dev/tty1"

function log_init()
{
  # 文件不存在则创建
  if [ ! -d "$log_file_dir" ]; then
    mkdir -p $log_file_dir
    chmod -R 755 $log_file_dir
  fi
  if [ ! -f "$log_file_path" ]; then
    touch $log_file_path
  fi
}

# 定义日志打印格式
function log()
{
  # 定义日志时间格式
  local curtime=$(date "+%Y-%m-%d %H:%M:%S")
  echo "$curtime $log_prefix $@"
  if [ "$user" == "root" ]; then
    echo "$curtime $log_prefix $@" >>$log_file_path
    #sleep 0.1
    #echo "$log_prefix [kmsg] $@" >$log_kmsg_path
    #sleep 0.1
    #echo "$curtime $log_prefix [con] $@" >$log_con_path
    #sleep 0.1
  fi
}

function log_from_file()
{
  tmpfile=$@
  if [ ! -f "$tmpfile" ]; then
    log "no file [$tmpfile] to log"
    return 1
  fi

  n=0
  while read line;
  do
    lists[$n]="$line"
    ((n++))
  done <$tmpfile

  # 定义日志时间格式
  local curtime=$(date "+%Y-%m-%d %H:%M:%S")

  #echo "$curtime $log_prefix $@"
  for i in $( seq 0 $((${#lists[*]} - 1)) )
  do
    echo "$curtime $log_prefix ${lists[i]}"
  done

  if [ "$user" == "root" ]; then
    #echo "$curtime $log_prefix $@" >>$log_file_path
    for i in $( seq 0 $((${#lists[*]} - 1)) )
    do
      echo "$curtime $log_prefix ${lists[i]}" >>$log_file_path
    done

    #sleep 0.1

    #echo "$log_prefix [kmsg] $@" >$log_kmsg_path
    #for i in $( seq 0 $((${#lists[*]} - 1)) )
    #do
    #  echo "$curtime $log_prefix ${lists[i]}" >$log_kmsg_path
    #  [ $((i%10)) -ne 9 ] || sleep 0.1
    #done

    #sleep 0.1

    #echo "$curtime $log_prefix [con] $@" >$log_con_path
    #for i in $( seq 0 $((${#lists[*]} - 1)) )
    #do
    #  echo "$curtime $log_prefix ${lists[i]}" >$log_con_path
    #done

    #sleep 0.1
  fi
}

function get_work_dir()
{
  ################
  # /kpcos/distro/yandi/releases/1.0/ports/amd64/host/yandi-host-setup-nocloud-arm64-image.sh
  ################

  local real_path=$(realpath $0)
  GLB_WORK_DIR=$(dirname "$real_path")
  log "WORK DIR is [$GLB_WORK_DIR]"
}

function build_image()
{
  img_size=$@
  if [ ${#img_size} -lt 2 ]; then
    log "length of input is [${#img_size}], too small!"
    return 1
  fi

  final=${img_size: -1}
  case "$final" in
    "g" ) ;;
    "t" ) ;;
    * )
      log "[$final] unit is incorrect"
      log "image size must be end with [g], [t] as for Giga, Tera"
      return 1
      ;;
  esac

  start=${img_size:0:1}
  case "$start" in
    [1-9] ) ;;
    * )
      log "[$start] as begin charactor is incorrect"
      log "begin char must be digit of [1-9]"
      return 1
      ;;
  esac

  middle=$(echo $img_size | awk '{print substr($0,2,length($0)-2)}')
  case "$middle" in
    [0-9] ) ;;
    "" ) ;;
    * )
      log "[$middle] is not digit [0-9] or empty string"
      return 1
      ;;
  esac

  # set std_val as in (MB) unit
  val=$(echo $img_size | awk '{print substr($0,1,length($0)-1)}')
  if [ "$final" == "t" ]; then
    std_val=$((val * 1024 * 1024))
  else
    std_val=$((val * 1024))
  fi
  std_size="${std_val}m"
  log "std_size=[$std_size]"

  real_size=$(echo "$std_size" | awk '{print tolower($0)}')
  log "real_size=[$real_size]"

  # 2. copy image file
  tmp_dl_img_file="tmp-${dl_img_file}"
  #cp -f debian-12-nocloud-arm64-20250416-2084.qcow2 tmp-debian-12-nocloud-arm64-20250416-2084.qcow2
  cp -f $dl_img_file $tmp_dl_img_file

  tmp_val=$(qemu-img info $tmp_dl_img_file | grep "virtual size:" | awk '{print substr($5,2,length($5)-1)}')
  log "tmp_val=[$tmp_val] Bytes"
  tmp_val=$((tmp_val / (1024 * 1024)))
  log "tmp_val=[$tmp_val] in MiB"

  # 3. resize image
  #qemu-img resize -f qcow2 tmp-debian-12-nocloud-arm64-20250416-2084.qcow2 40G

  if [ $std_val -lt $tmp_val ]; then
    log "need to shrink tmp image file"

    nbd_dev=
    for x in /sys/class/block/nbd* ; do
      S=$(cat $x/size)
      if [ "$S" == "0" ] ; then
        nbd_dev="/dev/$(basename $x)"
        break
      fi
    done
    if [ x"$nbd_dev" == x ]; then
      log "can't find free NBD device"
      return 1
    fi

    log "connect NBD $nbd_dev with $tmp_dl_img_file"
    qemu-nbd --connect=$nbd_dev $tmp_dl_img_file

    let counter=0
    while [ ! -e "${nbd_dev}p1" ]
    do
      if [ $counter -gt 9 ]; then
        log "wait too long for ${nbd_dev}p1 being ready, exit now."
	qemu-nbd --disconnect $nbd_dev
	exit 1
      fi
      let counter=$counter+1
      log "${nbd_dev}p1 is not ready, wait 1 second."
      sleep 1
    done

    # first resize ext4 filesystem, 128MiB is for p14, p15, preserve 8MiB for GPT
    log "resize2fs -f ${nbd_dev}p1 $((std_val-128-8))M"
    tmpfile=$(mktemp)
    /sbin/resize2fs -f ${nbd_dev}p1 "$((std_val-128-8))M" 2>&1 | tee $tmpfile
    log_from_file $tmpfile
    rm -f $tmpfile

    # preserve 4MiB for GPT
    log "parted $nbd_dev resizepart 1 $((std_val-4))MiB"
    tmpfile=$(mktemp)
    #parted $nbd_dev resizepart 1 "${std_val}MiB" 2>&1 | tee $tmpfile
    echo -e "Yes" | /sbin/parted ---pretend-input-tty "$nbd_dev" resizepart 1 "$((std_val-4))MiB" 2>&1 | tee $tmpfile 
    log_from_file $tmpfile
    rm -f $tmpfile

    log "disconnect NBD $nbd_dev"
    qemu-nbd --disconnect $nbd_dev
  fi

  #qemu-img resize --shrink $tmp_dl_img_file $real_size
  if [ "$std_val" != "$tmp_val" ]; then
    log "qemu-img resize --shrink $tmp_dl_img_file $std_size"
    qemu-img resize -f qcow2 --shrink $tmp_dl_img_file $std_size
  else
    log "no need to resize $tmp_dl_img_file $std_size"
  fi

  if [ $std_val -lt $tmp_val ]; then
    log "need to verify partitions on tmp image file"

    nbd_dev=
    for x in /sys/class/block/nbd* ; do
      S=$(cat $x/size)
      if [ "$S" == "0" ] ; then
        nbd_dev="/dev/$(basename $x)"
	nbd_dev_size_file="$x/size"
        break
      fi
    done
    if [ x"$nbd_dev" == x ]; then
      log "can't find free NBD device"
      return 1
    fi

    log "connect NBD $nbd_dev with $tmp_dl_img_file"
    qemu-nbd --connect=$nbd_dev $tmp_dl_img_file

    let counter=0
    while [ "$(cat $nbd_dev_size_file)" == "0" ]
    do
      if [ $counter -gt 9 ]; then
        log "wait too long for ${nbd_dev} being ready, exit now."
	qemu-nbd --disconnect $nbd_dev
	exit 1
      fi
      let counter=$counter+1
      log "${nbd_dev} is not ready, wait 1 second."
      sleep 1
    done

    log "verify partition of $nbd_dev"
    log "sgdisk -v $nbd_dev"
    tmpfile=$(mktemp)
    /sbin/sgdisk -v $nbd_dev 2>&1 | tee $tmpfile
    log_from_file $tmpfile
    rm -f $tmpfile

    log "fix partition error of $nbd_dev"
    log "sgdisk -e $nbd_dev"
    tmpfile=$(mktemp)
    /sbin/sgdisk -e $nbd_dev 2>&1 | tee $tmpfile
    log_from_file $tmpfile
    rm -f $tmpfile

    log "disconnect NBD $nbd_dev"
    qemu-nbd --disconnect $nbd_dev
  fi

  # 4. rename image
  #mv -f tmp-debian-12-nocloud-arm64-20250416-2084.qcow2 kpcos-yandi-1.02-nocloud-arm64-40g.qcow2
  target_img_file="kpcos-yandi-1.02-nocloud-arm64-${img_size}.qcow2"
  log "target image file name is [$target_img_file]"
  mv -f $tmp_dl_img_file $target_img_file

  # 5. show image file info
  tmpfile=$(mktemp)
  qemu-img info $target_img_file 2>&1 | tee $tmpfile
  log_from_file $tmpfile
  rm -f $tmpfile
}

# 0. init scripts
log_init
get_work_dir
cd "$GLB_WORK_DIR"

# 1. download image file from server
need_dl="yes"
#dl_img_file=debian-12-nocloud-arm64-20250416-2084.qcow2
dl_img_file=$GLB_DL_IMG_FILE
#dl_img_md5="666b76185ddb8d11d2053347f9ef3f65"
dl_img_md5=$GLB_DL_IMG_MD5
#dl_url=https://cloud.debian.org/images/cloud/bookworm/20250416-2084/$dl_img_file
dl_url=$GLB_DL_URL

if [ -f $dl_img_file ]; then
  MD5SUM=$(md5sum $dl_img_file | awk '{print $1}')
  if [ "${MD5SUM}" != "$dl_img_md5" ]; then
    log "File $dl_img_file is incorrect!"
    log "remove it and download again"
    rm -f $dl_img_file
    need_dl="yes"
  else
    log "File $dl_img_file is correct"
    need_dl="no"
  fi
fi
if [ "$need_dl" != "no" ]; then
  log "curl -o $dl_img_file -L $dl_url"
  curl -o $dl_img_file -L $dl_url
fi

if [ ! -f $dl_img_file ]; then
  echo "No file $dl_img_file"
  exit 1
fi
MD5SUM=$(md5sum $dl_img_file | awk '{print $1}')
if [ "${MD5SUM}" != "$dl_img_md5" ]; then
  echo "File $dl_img_file is incorrect"
  exit 1
fi

# 1. Enable NBD on the host
/sbin/modprobe nbd max_part=8

size_list=(
  "3g" "40g"
)

for size in "${size_list[@]}"
do
  log "target image size=[$size]"
  build_image $size
done

exit 0
